[None][feat] Kimi K3 NVFP4 on Rubin (SM107): Modeling V2 target and the runtime it needs - #21996

Leave 1 GiB outside the caching allocator cap for driver and KV transfer allocations created after executor setup. Cover the resulting allocator fraction with a CPU regression test.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…ention

With an FP8 KV cache, the FP8 context MLA FMHA reads FP8 Q, K and V, and the
attention op copies BF16 inputs to FP8 on every call. Let the caller hand it
E4M3 K / V directly, and optionally an FP8 Q through quant_q_buffer with
quant_scale_qkv, so that a caller that already produces FP8 tensors skips that
copy. BF16 inputs keep the existing path. Expose has_fp8_context_mla on
TrtllmAttention so that callers can tell when the op runs this FMHA.

Cover bit-identical outputs for FP8 and BF16 inputs on the cached-KV context
path, and the dtype check.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

Append drafter layers to the target cache manager and propagate their cache
layout through disaggregated transfer. Preserve token-parity lifecycle windows
and per-rank recurrent state, and validate supported page sizes. Drafters
declare what they read through DFlashForCausalLM hooks
(reads_managed_ctx_pages, reads_helix_sharded_ctx_pages, shift_label), so a
helix-sharded context is shared only with a drafter that reads it. Include
cache-manager, peer-metadata and speculative-worker regression coverage.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…ybrid layout

With the FP4 MLA layout, the hybrid manager scans and zeroes the KV cache at
warmup through _iter_fp4_mla_physical_pool_views, which listed every
physical pool except the one holding the appended draft layers. Yield that
single-buffer pool too, viewed with the draft layers' dtype and page shape,
so check_invalid_values_in_kv_cache(fill_with_zero=True) covers it.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…ing in the docs; test the reverse dtype mix and the draft bytes in the hybrid estimate

The speculative-decoding guide said a Kimi K3 hybrid target keeps a
standalone DSpark MLA drafter's layers with the CUTEDSL, TRTLLM or VANILLA
drafter backend. Whether a drafter shares is decided by its
reads_managed_ctx_pages declaration, and the default declaration covers only
the TRTLLM backend at trtllm-gen page sizes, so the guide now says that; the
list of separate-manager cases names the helix and opt-in exceptions.

The FP8 MLA context test also rejects BF16 K with FP8 V. A CPU test checks
that MambaHybridCacheManagerV2.get_cache_size_per_token adds the appended
draft layers' bytes per token to the slope and does not pass layer_specs or
shared_draft_block_width on to the hybrid estimate.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

An embedded DSpark draft (DeepSeek-V4 / V4.1) is a DSv4DSparkForCausalLM,
not a DFlashForCausalLM, and has no reads_managed_ctx_pages. The shared
draft KV gate read model.draft_model for every DSpark run and called that
hook, so creating the KV cache raised AttributeError where main reported
that the drafter has no paged block attention.

_loaded_standalone_dspark_drafter returns None for an embedded draft, and
_dflash_unify_blocker returns its reason before reading any hook when the
drafter has no block attention backend.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…sided config, and the receiver rejects a draft layout mismatch

With TRTLLM_DISAGG_SHARED_DRAFT_KV unset, disaggregated_mpi_worker set it to
"auto" for a config that lists only context servers (a DWDP context MPI
world), where main set "0". A generation worker launched separately with the
variable unset keeps a separate draft cache, so the two sides disagree.
_shared_draft_kv_default turns sharing on only when the config lists both
sides and passes the check; a one-sided config shares only when the variable
is "1".

The draft layer count and attention TP checks ran only in the sender's
register(), whose errors the listener thread logs, so a receiver whose draft
layout disagreed waited out kv_transfer_timeout_ms. They move to
PeerRegistrar.validate_draft_layers_compatible, which the receiver also runs
before it registers with a sender, so its requests to that sender fail with
PeerIncompatibleError.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…h holds

The cap is set after the final KV cache and CUDA graphs, so the 1 GiB kept
outside it comes out of the memory left free at that point. With less than
1 GiB free (a high free_gpu_memory_fraction on a small GPU) the cap dropped
below memory_reserved(), and with free + reserved under 1 GiB the fraction
was negative and set_per_process_memory_fraction raised ValueError. The
margin is now min(1 GiB, free memory).

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

Serialize DWDP worker URLs as YAML strings and publish every serving role to the same discovery directory when the launcher removes Slurm variables. This lets the frontend discover context workers and avoids waiting on a different directory.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…rd at startup

With num_serve_frontends > 1, a secondary frontend forwards the disaggregated
server's POST /steady_clock_offset to the launcher frontend over a Unix
socket. While the launcher frontend is not listening yet the forward is
refused: the frontend logs a warning that names ClientConnectorError and
ConnectionRefusedError and answers 503, and the disaggregated server leaves
that server's perf-metric timestamps unadjusted. The server log scan counted
the warning as a server error and failed the run. Treat exactly that warning
as benign; a refused forward of any other route, a timeout, or a line that
also carries another error keyword is still reported.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

… collectives

The kimi_k3_nvfp4__sm_107__default Modeling V2 target's deployment features, the
family routing to it, its core layers (linear, MXFP8 GEMM, norm, MLP, embedding,
all-reduce, streams), the Lamport collectives and the attention-residual norm,
with their ops/ entries and tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

The target's MLA attention: context FMHA and split-K context over the FP4 / FP8 KV
cache, FP4 MLA decode with the packed reduce, the fused q path and the Helix tail,
with their ops/ entries and tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

The target's KDA mixer and its prefill, MTP-decode and fused-projection kernels, with
their ops/ entries and tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

The target's routed and shared experts (router, route tail, small-batch MoE, ULL
front, MoE tail) and its low-M / mid-M FP8 and MXFP8 GEMMs with their push
variants, with their ops/ entries and tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

The target's DSpark drafter and spec-decode tail (vocab-parallel head, Markov chain,
acceptance, draft KV), with their ops/ entries and tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

The target's model (modeling.py), its checkpoint loading (weights.py) and the
runtime modules they use, with the model-level tests.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

…lists

GSM8K accuracy of the target's disaggregated deployments, their perf-sanity
configurations, and the CI and QA test lists for the target's tests.

Both roles of every deployment store the BF16 shared and latent MLP
projections as FP8 (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1; without it the DWDP
context servers run out of memory), keep the KV cache of the last two MLA
layers in FP8 (TLLM_KIMI_K3_FP8_KV_LAYERS=91,92) and use expandable PyTorch
allocator segments. The context servers cap their token capacity (1048576
with DWDP over 16 ranks, 131072 for the low-latency deployment): GPU 0 of each
node also holds the CUDA contexts of the node's other MPI ranks, which leaves
the default pool too little room there.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>

… and TCP transports

The perf-sanity launcher clears UCX_TLS for clusters without a UCX rule, so
the context and generation workers of the Kimi K3 SM107 gates ran with every
UCX transport enabled. The KV cache pool sized during the startup memory
profiling then stays registered with UCX after it is released, and its
memory does not come back for the final KV cache allocation: on the dep16
generation server, the KV cache manager's shutdown returns 2.8 of the
pool's 35.9 GiB with UCX_TLS unset, and all of it with UCX_TLS limited to
cuda_ipc,cuda_copy,sm,self,tcp.

Set UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp in the shared worker environment
of the four gate configurations (the transports the GB300 rules in
tests/test_common/cluster_env.py pin), and drop the context servers'
max_tokens caps, which only made room for the memory that was not
returned.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>